package com.demo.hdz.mediacodecdemo;

import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.media.MediaMuxer;
import android.os.Build;
import android.view.Surface;

import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.Arrays;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * Created by hdz on 2018/4/1.
 */

public class AVCodec {
    private final String TAG = "AVCodec";

    private final String MIME_TYPE = "video/avc"; //H264的MIME类型
    private static final int FRAME_RATE = 60;

    private MediaCodec m_mediaCodec = null;
    private Surface mSurface = null;
    private Lock m_lock = new ReentrantLock();

    private int m_validFrame = 0;
    private int frame_cnt = 0;
    private final int TIMEOUT_US = 10000;

    private MediaCodec.BufferInfo m_bufferInfo = new MediaCodec.BufferInfo();
    private AtomicBoolean m_quit = new AtomicBoolean(false);
    private int m_videoTrackIndex = -1;
    private boolean m_bMuxerStarted = false;
    private MediaMuxer m_mediaMuxer = null;

    public boolean initDecode(int width, int height, Surface surface) {
        //创建解码器
        try {
            if (m_mediaCodec == null) {
                m_mediaCodec = MediaCodec.createDecoderByType(MIME_TYPE);
                if(Build.VERSION.SDK_INT >= Build.VERSION_CODES.JELLY_BEAN_MR2) {
                    Logger.d("mediaCodec name: " + m_mediaCodec.getCodecInfo().getName());
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
            return false;
        }

        m_mediaCodec.stop();

        MediaFormat format = MediaFormat.createVideoFormat(MIME_TYPE, width, height);
        format.setInteger(MediaFormat.KEY_FRAME_RATE, FRAME_RATE);
        format.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1); //关键帧间隔时间 单位s
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.M) {
            format.setInteger(MediaFormat.KEY_ROTATION, 90); //图像旋转90度
        }

        if (!surface.isValid()) {
            Logger.e(TAG, "Surface is invalid!" + surface);
            return false;
        }

        //配置解码器
        try {
            m_mediaCodec.configure(format, surface, null, 0);
        } catch (Exception e) {
            e.printStackTrace();
            return false;
        }

        m_mediaCodec.start();

        return true;
    }

    public boolean initEncode(int width, int height, int bit_rate, int frame_rate, int r_frame_interval, MediaMuxer mediaMuxer) {
        m_mediaMuxer = mediaMuxer;

        MediaFormat format = MediaFormat.createVideoFormat(MIME_TYPE, width, height);
        //COLOR_FormatSurface这里表明数据将是一个graphicbuffer元数据
        format.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatSurface);
        //设置码率，码率越大视频越清晰，相应的占用内存也要更大
        format.setInteger(MediaFormat.KEY_BIT_RATE, bit_rate);
        //设置帧率，通常这个值越高，视频会显得越流畅，一般默认设置成30，最低可以设置成24，不要低于这个值，低于24会明显卡顿
        format.setInteger(MediaFormat.KEY_FRAME_RATE, frame_rate);
        //设置两个关键帧的间隔
        format.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, r_frame_interval);

        try {
            //创建编码器
            m_mediaCodec = MediaCodec.createEncoderByType(MIME_TYPE);

            m_mediaCodec.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);

            //这一步非常关键，它设置的是MediaCodec的编码源，也就是说，要告诉m_Encoder解码哪些流。
            //很出乎大家的意料，MediaCodec并没有要求我们传一个流文件进去，而是要求我们指定一个surface
            //而这个surface，其实就是MediaProjection中用来展示屏幕采集数据的surface。获取MediaCodec的surface，
            //这个surface其实就是个入口，屏幕作为输入源就会进入这个入口，然后交给MediaCodec编码
            mSurface = m_mediaCodec.createInputSurface();

            m_mediaCodec.start();
        } catch (Exception e) {
            e.printStackTrace();
            return false;
        }
        return true;
    }
    public Surface getSurface() {
        return mSurface;
    }


    public void decodeFile(InputStream inputStream) {
        try {
            if (inputStream != null) {
                while (true) {
                    try {
                        byte[] headData = new byte[4];
                        int ret = inputStream.read(headData, 0, headData.length);
                        if (ret == headData.length) {
                            ByteBuffer byteBuffer = ByteBuffer.wrap(Arrays.copyOf(headData, 4));
                            byteBuffer.order(ByteOrder.LITTLE_ENDIAN);

                            int dataLen = byteBuffer.getInt(0);
                            Logger.d("AvcDecoder, dataLen: " + dataLen);

                            byte[] inData = new byte[dataLen];
                            ret = inputStream.read(inData, 0, inData.length);
                            if (ret == inData.length) {
                                decode(inData);
                            } else {
                                inputStream.close();
                                break;
                            }
                        } else {
                            inputStream.close();
                            break;
                        }
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                }
                decode(null);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    //format naul + data:  00 00 00 01 xxxxxx
    public void decode(byte[] input) {
        long start_time = System.currentTimeMillis();

        m_lock.lock();

        try {
            boolean inputSuccess = false;
            while (!inputSuccess) {

                ByteBuffer[] inputBuffers = m_mediaCodec.getInputBuffers();   //输入缓冲区

                int inputBufferIndex = m_mediaCodec.dequeueInputBuffer(-1); //第一帧的帧索引
                Logger.d("Frame index of the first frame: " + inputBufferIndex);

                if (inputBufferIndex >= 0) {
                    ByteBuffer inputBuffer = inputBuffers[inputBufferIndex];
                    inputBuffer.clear();
                    if (input != null) {
                        inputBuffer.put(input);
                        m_mediaCodec.queueInputBuffer(inputBufferIndex, 0, input.length, frame_cnt++ * 1000 * 1000 / FRAME_RATE, 0);
                        Logger.d("input decode frame" + frame_cnt);
                    } else {
                        m_mediaCodec.queueInputBuffer(inputBufferIndex, 0, 0, frame_cnt++ * 1000 * 1000 / FRAME_RATE, MediaCodec.BUFFER_FLAG_END_OF_STREAM);
                    }
                    inputSuccess = true;
                }

                MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
                int outputBufferIndex = m_mediaCodec.dequeueOutputBuffer(bufferInfo, 1000 * 100);

                do {
                    if (outputBufferIndex == MediaCodec.INFO_TRY_AGAIN_LATER) {
                        Logger.d("INFO_TRY_AGAIN_LATER");
                    } else if (outputBufferIndex == MediaCodec.INFO_OUTPUT_BUFFERS_CHANGED) {
                        //encodeOutputBuffers = mDecodeMediaCodec.getOutputBuffers();
                        Logger.d("INFO_OUTPUT_BUFFERS_CHANGED");
                    } else if (outputBufferIndex == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED) {
                        MediaFormat format = m_mediaCodec.getOutputFormat();
                        Logger.d("onOutputFormatChanged: " + format.getInteger(MediaFormat.KEY_COLOR_FORMAT)
                                + "size: " + format.getInteger(MediaFormat.KEY_WIDTH) + "x" + format.getInteger(MediaFormat.KEY_HEIGHT));
                        //mediaformat changed
                    } else if (outputBufferIndex < 0) {
                        //unexpected result from encoder.dequeueOutputBuffer
                        Logger.d("outputBufferIndex < 0");
                    } else {
                        m_mediaCodec.releaseOutputBuffer(outputBufferIndex, true);
                        outputBufferIndex = m_mediaCodec.dequeueOutputBuffer(bufferInfo, 0);

                        m_validFrame++;
                        Logger.d("output decode frame: " + m_validFrame);
                    }
                } while (outputBufferIndex >= 0);
            }

        } catch (Throwable t) {
            t.printStackTrace();
        }
        m_lock.unlock();

        Logger.d("decode a frame use time: " + (System.currentTimeMillis() - start_time) + "ms");
    }

    public void close() {
        if (m_mediaCodec != null) {
            try {
                m_mediaCodec.stop();
                m_mediaCodec.release();
                m_mediaCodec = null;
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        m_quit.set(true);
    }



    public void recording() {
        while (!m_quit.get()) {
            //dequeueOutputBuffer方法可以这么理解：
            //它会出列一个输出buffer(可以理解为一帧画面)，
            //返回值是这一帧画面的顺序位置(类似于数组的下标)
            //第二个参数是超时时间，如果超过了这个时间还没成功出列，
            //那么就会跳过这一帧，去出列下一帧，并返回NFO_TRY_AGAIN_LATER标志位
            int index = m_mediaCodec.dequeueOutputBuffer(m_bufferInfo, TIMEOUT_US);

            //当格式改变的时候需要重新设置格式，第一次开始的时候会返回这个值
            if (index == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED) {
                resetOutputFormat();
            } else if (index == MediaCodec.INFO_TRY_AGAIN_LATER) {
                try {
                    Thread.sleep(10);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            } else if (index > 0) { //到这里说明dequeueOutputBuffer执行正常
                if (!m_bMuxerStarted) {
                    continue;
                }

                //转换成mp4
                encodeToVideoTrack(index);

                ////释放缓存的资源
                m_mediaCodec.releaseOutputBuffer(index, false);
            }
        }
    }

    private void resetOutputFormat() {
        if (m_bMuxerStarted) {
            return;
        }
        //将MediaCodec的Format设置给MediaMuxer
        MediaFormat newFormat = m_mediaCodec.getOutputFormat();

        //获取m_videoTrackIndex，这个值是每一帧画面要放置的顺序
        m_videoTrackIndex = m_mediaMuxer.addTrack(newFormat);
        m_mediaMuxer.start();
        m_bMuxerStarted = true;
    }

    //这里是将数据传给MediaMuxer，将其转换成mp4
    private void encodeToVideoTrack(int index) {
        //通过index获取到ByteBuffer(可以理解为一帧)
        ByteBuffer encodeData = m_mediaCodec.getOutputBuffer(index);

        //当m_bufferInfo返回这个标志位时，就说明已经传完数据了
        //将m_bufferInfo.size设为0，准备将其回收
        if ((m_bufferInfo.flags & MediaCodec.BUFFER_FLAG_CODEC_CONFIG) !=0) {
            m_bufferInfo.size = 0;
        }
        if (m_bufferInfo.size == 0) {
            encodeData = null;
        } else {
            Logger.d("", "got buffer, info: size="+m_bufferInfo.size+",presentationTimeUs="+m_bufferInfo.presentationTimeUs+", offset="+m_bufferInfo.offset);
        }
        if (encodeData != null) {
            //设置该从哪个位置读取数据
            encodeData.position(m_bufferInfo.offset);
            //设置该读取多少数据
            encodeData.limit(m_bufferInfo.offset + m_bufferInfo.size);
            //将数据写入到文件
            //第一个参数是每一帧画面要放置的顺序
            //第二个参数是要写入的数据
            //第三个参数是BufferInfo，这个数据包含encodeData的offset和size
            m_mediaMuxer.writeSampleData(m_videoTrackIndex, encodeData, m_bufferInfo);
        }
    }


}
